import pandas as pd

from shannon import bestSplit, mySplit

row_data = {'是否陪伴': [0, 0, 0, 1, 1],
            '是否玩游戏': [1, 1, 0, 1, 1],
            '渣男': ['是', '是', '不是', '不是', '不是']}
dataSet = pd.DataFrame(row_data)
def createTree(dataSet):
    """
    函数功能：基于最大信息增益切分数据集，递归构建决策树
    参数说明：
    dataSet：原始数据集(最有一列是标签)
    return：myTree：字典形式的树
    """
    featlist = list(dataSet.columns)  # 提取出数据集所有的列
    # print(featlist) # ['是否陪伴', '是否玩游戏', '渣男']
    # 最后一列 -> [所有出现的取值:该取值的数量]
    classlist = dataSet.iloc[:, -1].value_counts()  # 获取最后一列类标签
    # 每个分支下的所有实例都具有相同的分类 || 程序遍历完所有划分数据集的属性
    # 判断最多标签数目是否等于数据集行数，或者数据集是否只有一列
    if classlist[0] == dataSet.shape[0] or dataSet.shape[1] == 1:
        return classlist.index[0]  # 如果是，返回类标签
    axis = bestSplit(dataSet)  # 确定出当前最佳切分列的索引
    bestfeat = featlist[axis]  # 获取该索引对应的特征
    myTree = {bestfeat: {}}  # 采用字典嵌套的方式存储树信息
    del featlist[axis]  # 删除当前特征
    valuelist = set(dataSet.iloc[:, axis])  # 提取最佳切分列所有属性值
    for value in valuelist:  # 对每一个属性值递归建树
        # mySplit([[1,1],[2,1],[0,0]],1,1) -> [1,2]
        # mySplit(数据集,第i列,value) 提取dataset中,第i列中,和value相等的值,组成新dataset
        myTree[bestfeat][value] = createTree(mySplit(dataSet, axis, value))
    return myTree

#查看运行结果
myTree = createTree(dataSet)
# {'是否陪伴': {0: {'是否玩游戏': {0: '不是', 1: '是'}}, 1: '不是'}}
print(myTree)